{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 解读类型推断"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import testing\n",
    "import pytest\n",
    "import tvm\n",
    "from tvm import IRModule, relay\n",
    "from tvm.relay import op, transform\n",
    "from tvm.relay.op import op as _op\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def infer_mod(mod, annotate_spans=True):\n",
    "    if annotate_spans:\n",
    "        mod = relay.transform.AnnotateSpans()(mod)\n",
    "\n",
    "    mod = transform.InferType()(mod)\n",
    "    return mod\n",
    "\n",
    "\n",
    "def infer_expr(expr):\n",
    "    transform.InferTypeLocal(expr)\n",
    "    return expr\n",
    "\n",
    "\n",
    "def assert_has_type(expr, typ, mod=None):\n",
    "    if not mod:\n",
    "        mod = tvm.IRModule({})\n",
    "\n",
    "    mod[\"main\"] = expr\n",
    "    mod = infer_mod(mod)\n",
    "    checked_expr = mod[\"main\"]\n",
    "    checked_type = checked_expr.checked_type\n",
    "    if checked_type != typ:\n",
    "        raise RuntimeError(\"Type mismatch %s vs %s\" % (checked_type, typ))\n",
    "\n",
    "\n",
    "def initialize_box_adt(mod):\n",
    "    # initializes simple ADT for tests\n",
    "    box = relay.GlobalTypeVar(\"box\")\n",
    "    tv = relay.TypeVar(\"tv\")\n",
    "    constructor = relay.Constructor(\"constructor\", [tv], box)\n",
    "    data = relay.TypeData(box, [tv], [constructor])\n",
    "    mod[box] = data\n",
    "    return box, constructor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试单态 let 类型推断\n",
    "\n",
    "monomorphic：单态\n",
    "\n",
    "程序：`let %x = 1; %x`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "sb = relay.ScopeBuilder()\n",
    "x = relay.var(\"x\", dtype=\"float64\", shape=())\n",
    "x = sb.let(x, relay.const(1.0, \"float64\"))\n",
    "sb.ret(x)\n",
    "out = sb.get()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "let %x: float64 = 1f64;\n",
      "%x\n"
     ]
    }
   ],
   "source": [
    "print(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "xchecked = infer_expr(out)\n",
    "assert xchecked.checked_type == relay.scalar_type(\"float64\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试单个算子类型推断"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = relay.var(\"x\", shape=[])\n",
    "func = relay.Function([x], op.log(x))\n",
    "ttype = relay.TensorType([], dtype=\"float32\")\n",
    "assert_has_type(func, relay.FuncType([ttype], ttype))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fn (%x: float32) {\n",
      "  log(%x)\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(func)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 其他类型推断"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def test_add_broadcast_op():\n",
    "    \"\"\"\n",
    "    Program:\n",
    "        fn (%x: Tensor[(10, 4), float32], %y: Tensor[(5, 10, 1), float32])\n",
    "            -> Tensor[(5, 10, 4), float32] {\n",
    "            %x + %y\n",
    "        }\n",
    "    \"\"\"\n",
    "    x = relay.var(\"x\", shape=(10, 4))\n",
    "    y = relay.var(\"y\", shape=(5, 10, 1))\n",
    "    z = x + y\n",
    "    func = relay.Function([x, y], z)\n",
    "    t1 = relay.TensorType((10, 4), \"float32\")\n",
    "    t2 = relay.TensorType((5, 10, 1), \"float32\")\n",
    "    t3 = relay.TensorType((5, 10, 4), \"float32\")\n",
    "    expected_ty = relay.FuncType([t1, t2], t3)\n",
    "    assert_has_type(func, expected_ty)\n",
    "\n",
    "\n",
    "def test_dual_op():\n",
    "    \"\"\"Program:\n",
    "    fn (%x : Tensor[(10, 10), float32]) {\n",
    "      let %t1 = log(x);\n",
    "      let %t2 = add(%t1, %x);\n",
    "      %t1\n",
    "    }\n",
    "    \"\"\"\n",
    "    tp = relay.TensorType((10, 10), \"float32\")\n",
    "    x = relay.var(\"x\", tp)\n",
    "    sb = relay.ScopeBuilder()\n",
    "    t1 = sb.let(\"t1\", relay.log(x))\n",
    "    t2 = sb.let(\"t2\", relay.add(t1, x))\n",
    "    sb.ret(t2)\n",
    "    f = relay.Function([x], sb.get())\n",
    "    fchecked = infer_expr(f)\n",
    "    assert fchecked.checked_type == relay.FuncType([tp], tp)\n",
    "\n",
    "\n",
    "def test_decl():\n",
    "    \"\"\"Program:\n",
    "    def @f(%x : Tensor[(10, 10), float32]) {\n",
    "        log(%x)\n",
    "    }\n",
    "    \"\"\"\n",
    "    tp = relay.TensorType((10, 10))\n",
    "    x = relay.var(\"x\", tp)\n",
    "    f = relay.Function([x], relay.log(x))\n",
    "    fchecked = infer_expr(f)\n",
    "    assert fchecked.checked_type == relay.FuncType([tp], tp)\n",
    "\n",
    "\n",
    "def test_recursion():\n",
    "    \"\"\"\n",
    "    Program:\n",
    "       def @f(%n: int32, %data: float32) -> float32 {\n",
    "          if (%n == 0) {\n",
    "              %data\n",
    "          } else {\n",
    "              @f(%n - 1, log(%data))\n",
    "          }\n",
    "       }\n",
    "    \"\"\"\n",
    "    sb = relay.ScopeBuilder()\n",
    "    f = relay.GlobalVar(\"f\")\n",
    "    ti32 = relay.scalar_type(\"int32\")\n",
    "    tf32 = relay.scalar_type(\"float32\")\n",
    "    n = relay.var(\"n\", ti32)\n",
    "    data = relay.var(\"data\", tf32)\n",
    "\n",
    "    with sb.if_scope(relay.equal(n, relay.const(0, ti32))):\n",
    "        sb.ret(data)\n",
    "    with sb.else_scope():\n",
    "        sb.ret(f(relay.subtract(n, relay.const(1, ti32)), relay.log(data)))\n",
    "    mod = tvm.IRModule()\n",
    "    mod[f] = relay.Function([n, data], sb.get())\n",
    "    mod = infer_mod(mod)\n",
    "    assert \"@f(%1, %2)\" in mod.astext()\n",
    "    assert mod[\"f\"].checked_type == relay.FuncType([ti32, tf32], tf32)\n",
    "\n",
    "\n",
    "def test_incomplete_call():\n",
    "    tt = relay.scalar_type(\"int32\")\n",
    "    x = relay.var(\"x\", tt)\n",
    "    f_type = relay.FuncType([tt], tt)\n",
    "    f = relay.var(\"f\")\n",
    "    func = relay.Function([x, f], relay.Call(f, [x]), tt)\n",
    "\n",
    "    ft = infer_expr(func)\n",
    "    assert ft.checked_type == relay.FuncType([tt, f_type], tt)\n",
    "\n",
    "\n",
    "def test_higher_order_argument():\n",
    "    a = relay.TypeVar(\"a\")\n",
    "    x = relay.Var(\"x\", a)\n",
    "    id_func = relay.Function([x], x, a, [a])\n",
    "\n",
    "    b = relay.TypeVar(\"b\")\n",
    "    f = relay.Var(\"f\", relay.FuncType([b], b))\n",
    "    y = relay.Var(\"y\", b)\n",
    "    ho_func = relay.Function([f, y], f(y), b, [b])\n",
    "\n",
    "    # id func should be an acceptable argument to the higher-order\n",
    "    # function even though id_func takes a type parameter\n",
    "    ho_call = ho_func(id_func, relay.const(0, \"int32\"))\n",
    "\n",
    "    hc = infer_expr(ho_call)\n",
    "    expected = relay.scalar_type(\"int32\")\n",
    "    assert hc.checked_type == expected\n",
    "\n",
    "\n",
    "def test_higher_order_return():\n",
    "    a = relay.TypeVar(\"a\")\n",
    "    x = relay.Var(\"x\", a)\n",
    "    id_func = relay.Function([x], x, a, [a])\n",
    "\n",
    "    b = relay.TypeVar(\"b\")\n",
    "    nested_id = relay.Function([], id_func, relay.FuncType([b], b), [b])\n",
    "\n",
    "    ft = infer_expr(nested_id)\n",
    "    assert ft.checked_type == relay.FuncType([], relay.FuncType([b], b), [b])\n",
    "\n",
    "\n",
    "def test_higher_order_nested():\n",
    "    a = relay.TypeVar(\"a\")\n",
    "    x = relay.Var(\"x\", a)\n",
    "    id_func = relay.Function([x], x, a, [a])\n",
    "\n",
    "    choice_t = relay.FuncType([], relay.scalar_type(\"bool\"))\n",
    "    f = relay.Var(\"f\", choice_t)\n",
    "\n",
    "    b = relay.TypeVar(\"b\")\n",
    "    z = relay.Var(\"z\")\n",
    "    top = relay.Function(\n",
    "        [f], relay.If(f(), id_func, relay.Function([z], z)), relay.FuncType([b], b), [b]\n",
    "    )\n",
    "\n",
    "    expected = relay.FuncType([choice_t], relay.FuncType([b], b), [b])\n",
    "    ft = infer_expr(top)\n",
    "    assert ft.checked_type == expected\n",
    "\n",
    "\n",
    "def test_tuple():\n",
    "    tp = relay.TensorType((10,))\n",
    "    x = relay.var(\"x\", tp)\n",
    "    res = relay.Tuple([x, x])\n",
    "    assert infer_expr(res).checked_type == relay.TupleType([tp, tp])\n",
    "\n",
    "\n",
    "def test_ref():\n",
    "    x = relay.var(\"x\", \"float32\")\n",
    "    y = relay.var(\"y\", \"float32\")\n",
    "    r = relay.RefCreate(x)\n",
    "    st = relay.scalar_type(\"float32\")\n",
    "    assert infer_expr(r).checked_type == relay.RefType(st)\n",
    "    g = relay.RefRead(r)\n",
    "    assert infer_expr(g).checked_type == st\n",
    "    w = relay.RefWrite(r, y)\n",
    "    assert infer_expr(w).checked_type == relay.TupleType([])\n",
    "\n",
    "\n",
    "def test_free_expr():\n",
    "    x = relay.var(\"x\", \"float32\")\n",
    "    y = relay.add(x, x)\n",
    "    yy = infer_expr(y)\n",
    "    assert tvm.ir.structural_equal(yy.args[0], x, map_free_vars=True)\n",
    "    assert yy.checked_type == relay.scalar_type(\"float32\")\n",
    "    assert x.vid.same_as(yy.args[0].vid)\n",
    "\n",
    "\n",
    "def test_type_args():\n",
    "    x = relay.var(\"x\", shape=(10, 10))\n",
    "    y = relay.var(\"y\", shape=(1, 10))\n",
    "    z = relay.add(x, y)\n",
    "\n",
    "    # InferTypeLocal does not support populating the type_args field\n",
    "    mod = infer_mod(IRModule.from_expr(z))\n",
    "    mod = infer_mod(mod, annotate_spans=False)\n",
    "    ty_args = mod[\"main\"].body.type_args\n",
    "    assert len(ty_args) == 2\n",
    "    assert ty_args[0].dtype == \"float32\"\n",
    "    assert ty_args[1].dtype == \"float32\"\n",
    "    sh1 = ty_args[0].shape\n",
    "    sh2 = ty_args[1].shape\n",
    "    assert sh1[0].value == 10\n",
    "    assert sh1[1].value == 10\n",
    "    assert sh2[0].value == 1\n",
    "    assert sh2[1].value == 10\n",
    "\n",
    "\n",
    "def test_global_var_recursion():\n",
    "    mod = tvm.IRModule({})\n",
    "    gv = relay.GlobalVar(\"main\")\n",
    "    x = relay.var(\"x\", shape=[])\n",
    "    tt = relay.scalar_type(\"float32\")\n",
    "\n",
    "    func = relay.Function([x], relay.Call(gv, [x]), tt)\n",
    "    mod[gv] = func\n",
    "    mod = infer_mod(mod)\n",
    "    func_ty = mod[\"main\"].checked_type\n",
    "\n",
    "    assert func_ty == relay.FuncType([tt], tt)\n",
    "\n",
    "\n",
    "def test_equal():\n",
    "    i = relay.var(\"i\", shape=[], dtype=\"int32\")\n",
    "    eq = op.equal(i, relay.const(0, dtype=\"int32\"))\n",
    "    func = relay.Function([i], eq)\n",
    "    ft = infer_expr(func)\n",
    "    expected = relay.FuncType([relay.scalar_type(\"int32\")], relay.scalar_type(\"bool\"))\n",
    "    assert ft.checked_type == expected\n",
    "\n",
    "    assert ft.checked_type == relay.FuncType(\n",
    "        [relay.scalar_type(\"int32\")], relay.scalar_type(\"bool\")\n",
    "    )\n",
    "\n",
    "\n",
    "def test_constructor_type():\n",
    "    mod = tvm.IRModule()\n",
    "    box, constructor = initialize_box_adt(mod)\n",
    "\n",
    "    a = relay.TypeVar(\"a\")\n",
    "    x = relay.Var(\"x\", a)\n",
    "    func = relay.Function([x], constructor(x), box(a), [a])\n",
    "    mod[\"main\"] = func\n",
    "    mod = infer_mod(mod)\n",
    "    func_ty = mod[\"main\"].checked_type\n",
    "    box = mod.get_global_type_var(\"box\")\n",
    "    expected = relay.FuncType([a], box(a), [a])\n",
    "    assert func_ty == expected\n",
    "\n",
    "\n",
    "def test_constructor_call():\n",
    "    mod = tvm.IRModule()\n",
    "    box, constructor = initialize_box_adt(mod)\n",
    "\n",
    "    box_unit = constructor(relay.Tuple([]))\n",
    "    box_constant = constructor(relay.const(0, \"float32\"))\n",
    "\n",
    "    func = relay.Function([], relay.Tuple([box_unit, box_constant]))\n",
    "    mod[\"main\"] = func\n",
    "    mod = infer_mod(mod)\n",
    "    ret_type = mod[\"main\"].checked_type.ret_type.fields\n",
    "    # NB(@jroesch): when we annotate spans the ast fragments before\n",
    "    # annotation the previous fragments will no longer be directly equal.\n",
    "    box = mod.get_global_type_var(\"box\")\n",
    "    expected1 = box(relay.TupleType([]))\n",
    "    expected2 = box(relay.TensorType((), \"float32\"))\n",
    "    assert ret_type[0] == expected1\n",
    "    assert ret_type[1] == expected2\n",
    "\n",
    "\n",
    "def test_adt_match():\n",
    "    mod = tvm.IRModule()\n",
    "    box, constructor = initialize_box_adt(mod)\n",
    "\n",
    "    v = relay.Var(\"v\", relay.TensorType((), \"float32\"))\n",
    "    match = relay.Match(\n",
    "        constructor(relay.const(0, \"float32\")),\n",
    "        [\n",
    "            relay.Clause(\n",
    "                relay.PatternConstructor(constructor, [relay.PatternVar(v)]), relay.Tuple([])\n",
    "            ),\n",
    "            # redundant but shouldn't matter to typechecking\n",
    "            relay.Clause(relay.PatternWildcard(), relay.Tuple([])),\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    func = relay.Function([], match)\n",
    "    mod[\"main\"] = func\n",
    "    mod = infer_mod(mod)\n",
    "    actual = mod[\"main\"].checked_type.ret_type\n",
    "    assert actual == relay.TupleType([])\n",
    "\n",
    "\n",
    "def test_adt_match_type_annotations():\n",
    "    mod = tvm.IRModule()\n",
    "    box, constructor = initialize_box_adt(mod)\n",
    "\n",
    "    # the only type annotation is inside the match pattern var\n",
    "    # but that should be enough info\n",
    "    tt = relay.TensorType((2, 2), \"float32\")\n",
    "    x = relay.Var(\"x\")\n",
    "    mv = relay.Var(\"mv\", tt)\n",
    "    match = relay.Match(\n",
    "        constructor(x),\n",
    "        [\n",
    "            relay.Clause(\n",
    "                relay.PatternConstructor(constructor, [relay.PatternVar(mv)]), relay.Tuple([])\n",
    "            )\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    mod[\"main\"] = relay.Function([x], match)\n",
    "    mod = infer_mod(mod)\n",
    "    ft = mod[\"main\"].checked_type\n",
    "    assert ft == relay.FuncType([tt], relay.TupleType([]))\n",
    "\n",
    "\n",
    "def test_let_polymorphism():\n",
    "    id = relay.Var(\"id\")\n",
    "    xt = relay.TypeVar(\"xt\")\n",
    "    x = relay.Var(\"x\", xt)\n",
    "    body = relay.Tuple([id(relay.const(1)), id(relay.Tuple([]))])\n",
    "    body = relay.Let(id, relay.Function([x], x, xt, [xt]), body)\n",
    "    body = infer_expr(body)\n",
    "    int32 = relay.TensorType((), \"int32\")\n",
    "    tvm.ir.assert_structural_equal(body.checked_type, relay.TupleType([int32, relay.TupleType([])]))\n",
    "\n",
    "\n",
    "def test_type_arg_infer():\n",
    "    code = \"\"\"\n",
    "#[version = \"0.0.5\"]\n",
    "def @id[A](%x: A) -> A {\n",
    "  %x\n",
    "}\n",
    "def @main(%f: float32) -> float32 {\n",
    "  @id(%f)\n",
    "}\n",
    "\"\"\"\n",
    "    mod = tvm.relay.fromtext(code)\n",
    "    mod = transform.InferType()(mod)\n",
    "    tvm.ir.assert_structural_equal(mod[\"main\"].body.type_args, [relay.TensorType((), \"float32\")])\n",
    "\n",
    "\n",
    "def test_dynamic_function():\n",
    "    dy_tt = relay.TensorType([relay.Any()], \"float32\")\n",
    "    s_tt = relay.TensorType([10], \"float32\")\n",
    "    x = relay.Var(\"x\", dy_tt)\n",
    "    f = relay.Function([x], x + x)\n",
    "    y = relay.Var(\"y\", s_tt)\n",
    "    c = f(y)\n",
    "\n",
    "    mod = tvm.IRModule()\n",
    "    mod[\"main\"] = relay.Function([y], c)\n",
    "    mod = transform.InferType()(mod)\n",
    "    assert mod[\"main\"].params[0].checked_type == s_tt\n",
    "\n",
    "    data = relay.var(\n",
    "        \"data\", shape=(relay.Any(), relay.Any(), relay.Any(), relay.Any()), dtype=\"float32\"\n",
    "    )\n",
    "    weigth = relay.const(np.full((16, 16, 3, 3), 0.25), dtype=\"float32\")\n",
    "    x = relay.nn.conv2d(data, weigth, kernel_size=(3, 3), channels=16, groups=2)\n",
    "    mod = tvm.IRModule.from_expr(x)\n",
    "    mod = transform.InferType()(mod)\n",
    "\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    tvm.testing.main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试参数约简以推断返回类型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_shape = (1, 1)\n",
    "broadcast_shape = [1, 1]\n",
    "shape_dtypes = [(\"int32\", lambda x: np.int32(x)), (\"int64\", lambda x: np.int64(x))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Testing with argmax （argreduce）\n",
    "for (sdtype, conv) in shape_dtypes:\n",
    "    x = relay.var(\"data\", relay.TensorType(x_shape, \"float32\"))\n",
    "    broadcast_to = relay.op.broadcast_to(x, relay.const(broadcast_shape, dtype=sdtype))\n",
    "    argmax = relay.op.argmax(broadcast_to, axis=[1])\n",
    "\n",
    "    f = relay.Function([x], argmax)\n",
    "    assert_has_type(\n",
    "        f,\n",
    "        relay.FuncType(\n",
    "            [relay.TensorType(broadcast_shape, \"float32\")],\n",
    "            relay.TensorType([conv(1)], dtype=sdtype),\n",
    "        ),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Testing with argmin\n",
    "for (sdtype, conv) in shape_dtypes:\n",
    "    x = relay.var(\"data\", relay.TensorType(x_shape, \"float32\"))\n",
    "    broadcast_to = relay.op.broadcast_to(x, relay.const(broadcast_shape, dtype=sdtype))\n",
    "    argmin = relay.op.argmin(broadcast_to, axis=[1])\n",
    "\n",
    "    f = relay.Function([x], argmin)\n",
    "    assert_has_type(\n",
    "        f,\n",
    "        relay.FuncType(\n",
    "            [relay.TensorType(broadcast_shape, \"float32\")],\n",
    "            relay.TensorType([conv(1)], dtype=sdtype),\n",
    "        ),\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py311",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
